package com.beta.cat.config;

import org.apache.commons.lang3.StringUtils;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.config.spring.extension.SpringExtensionFactory;
import org.apache.dubbo.rpc.*;
import org.slf4j.MDC;
import org.springframework.context.ApplicationContext;
import org.springframework.core.env.Environment;

import java.util.UUID;

/**
 * 日志染色ProviderRpcTraceFilter
 *
 * @author phpdragon
 */
@Activate(group = {CommonConstants.PROVIDER, CommonConstants.CONSUMER}, order = 1)
public class RpcTraceFilter implements Filter {
    /**
     * MDC里面的跟踪key名称
     */
    public static final String MDC_TRACE_KEY = "traceId";
    /**
     * MDC里面的应用名称key名称
     */
    public static final String MDC_APPNAME_KEY = "appName";
    private Environment environment;
    private ITraceMDC mdc;

    public Environment getEnvironment() {
        return environment;
    }

    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }

    public ITraceMDC getMdc() {
        if (mdc == null) {
            for (ApplicationContext ctx : SpringExtensionFactory.getContexts()) {
                try {
                    mdc = ctx.getBean(ITraceMDC.class);
                    break;
                } catch (Exception e) {
                }
            }
            if (mdc == null) {
                mdc = new ITraceMDC() {
                    @Override
                    public void put(String key, String val) {
                    }

                    @Override
                    public String get(String key) {
                        return "";
                    }
                };
            }
        }
        return mdc;
    }

    public void setMdc(ITraceMDC mdc) {
        this.mdc = mdc;
    }

    public String getAppName() {
        return getEnvironment().getProperty("spring.application.name", "notFound");
    }

    /**
     * @param invoker
     * @param invocation
     * @return
     * @throws RpcException
     */
    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        String traceId = RpcContext.getContext().getAttachment(MDC_TRACE_KEY);
        if (StringUtils.isBlank(traceId)) {
            traceId = MDC.get(MDC_TRACE_KEY);
        }
        if (StringUtils.isBlank(traceId)) {
            traceId = getUUID();
        }
//        MDC.put("traceId", traceId);
        getMdc().put(MDC_TRACE_KEY, traceId);
        getMdc().put(MDC_APPNAME_KEY, getAppName());

        RpcContext.getContext().setAttachment(MDC_TRACE_KEY, traceId);
        long start = System.currentTimeMillis();
        try {
            return invoker.invoke(invocation);
        } finally {
            RpcContext.getContext().setAttachment(MDC_TRACE_KEY, null);
        }
    }

    /**
     * 获取UUID
     *
     * @return String UUID
     */
    public static String getUUID() {
        String uuid = UUID.randomUUID().toString();
        //替换-字符
        return uuid.replaceAll("-", "");
    }
}